import psycopg2
import re

# 读取redshift系统表，导出建表语句
def main():
    # Redshift连接配置
    HOST = '127.0.0.1'
    PORT = 54391
    DATABASE = 'bigdata'
    USER = 'root'
    PASSWORD = 'root'

    try:
        # 建立数据库连接
        conn = psycopg2.connect(
            host=HOST,
            port=PORT,
            dbname=DATABASE,
            user=USER,
            password=PASSWORD
        )
    except Exception as e:
        print(f"数据库连接失败: {e}")
        return

    # 创建游标
    cur = conn.cursor()

    try:
        # 获取 bigdata.dwd 模式下的所有表名
        cur.execute("""
            SELECT table_name
            FROM information_schema.tables
            WHERE table_schema = 'dwd'
              AND table_catalog = 'bigdata'
              AND table_type = 'BASE TABLE'
            ORDER BY table_name;  -- 为了一致性，按表名排序
        """)

        tables = cur.fetchall()
    except Exception as e:
        print(f"查询表名失败: {e}")
        cur.close()
        conn.close()
        return

    # 准备一个列表来保存 CREATE TABLE 语句
    ddl_statements = []
    table_count = 0
    file_count = 0  # 修改初始值，使文件名不带数字

    for table in tables:
        table_name = table[0]
        try:
            # 执行 'show table bigdata.dwd.table_name;'
            show_table_query = f"show table bigdata.dwd.\"{table_name}\";"
            cur.execute(show_table_query)

            # 获取 CREATE TABLE 语句
            create_table_stmt = ''
            for row in cur.fetchall():
                create_table_stmt += row[0] + '\n'

            # **修改 CREATE TABLE 语句**

            # 1. 将 'dwd' 模式替换为 'bigdata.yujiahao' 模式
            # 使用正则表达式仅替换 CREATE TABLE 行中的模式
            create_table_stmt = re.sub(
                r'CREATE TABLE dwd\.(\S+)',
                r'CREATE TABLE bigdata.yujiahao.\1',
                create_table_stmt,
                flags=re.IGNORECASE
            )

            # 2. 用双引号括起所有列名，但不包括约束

            # 定义一个函数来分割列定义，考虑括号嵌套
            def split_columns(columns_str):
                columns = []
                bracket_level = 0
                current_col = ''
                i = 0
                while i < len(columns_str):
                    char = columns_str[i]
                    if char == ',' and bracket_level == 0:
                        columns.append(current_col.strip())
                        current_col = ''
                    else:
                        if char == '(':
                            bracket_level += 1
                        elif char == ')':
                            bracket_level -= 1
                        current_col += char
                    i += 1
                if current_col.strip():
                    columns.append(current_col.strip())
                return columns

            # 定义一个函数来获取列定义部分
            def get_columns_section(stmt):
                start_idx = stmt.find('(')
                if start_idx == -1:
                    return None, stmt, ''
                idx = start_idx + 1
                bracket_level = 1
                while idx < len(stmt) and bracket_level > 0:
                    if stmt[idx] == '(':
                        bracket_level += 1
                    elif stmt[idx] == ')':
                        bracket_level -= 1
                    idx += 1
                if bracket_level == 0:
                    columns_section = stmt[start_idx+1:idx-1]  # 不包括最外层括号
                    before_columns = stmt[:start_idx+1]        # 包含 '('
                    after_columns = stmt[idx-1:]               # 包含 ')'
                    return columns_section, before_columns, after_columns
                else:
                    return None, stmt, ''

            # 获取列定义部分
            columns_section, before_columns, after_columns = get_columns_section(create_table_stmt)
            if columns_section:
                # 分割列定义
                column_defs = split_columns(columns_section)
                new_column_defs = []
                for col_def in column_defs:
                    # 检查是否是约束
                    constraint_keywords = ['PRIMARY KEY', 'FOREIGN KEY', 'UNIQUE', 'CHECK', 'CONSTRAINT', 'REFERENCES']
                    is_constraint = False
                    for keyword in constraint_keywords:
                        if col_def.strip().upper().startswith(keyword):
                            is_constraint = True
                            break
                    if is_constraint:
                        # 不修改约束，直接添加
                        new_column_defs.append(col_def.strip())
                    else:
                        # 匹配列名和剩余部分
                        col_match = re.match(r'\s*(\S+)\s+(.*)', col_def.strip(), re.DOTALL)
                        if col_match:
                            column_name = col_match.group(1)
                            rest = col_match.group(2)

                            # 在检查列名时，避免将 NOT NULL、DEFAULT 等作为列名的一部分
                            # 处理列级别的约束
                            tokens = rest.strip().upper().split()
                            if tokens and tokens[0] in ['NOT', 'NULL', 'DEFAULT', 'ENCODE']:
                                # 这可能是列名和约束之间没有空格，处理可能的异常
                                pass

                            # 检查列名是否已被引号括起
                            if not (column_name.startswith('"') and column_name.endswith('"')):
                                # 用双引号括起列名
                                column_name = f'"{column_name}"'

                            new_col_def = f"{column_name} {rest}".strip()
                            new_column_defs.append(new_col_def)
                        else:
                            # 如果无法匹配，保留原始定义
                            new_column_defs.append(col_def.strip())

                # 重新组合列定义部分
                modified_columns_section = '\n  ' + ',\n  '.join(new_column_defs) + '\n'
                # 重建 CREATE TABLE 语句
                create_table_stmt = before_columns + modified_columns_section + after_columns
            else:
                # 如果未找到列定义部分，保持原样
                pass

            # 3. 如有必要，调整 DISTSTYLE 和 DISTKEY

            # 检查是否指定了 DISTSTYLE AUTO
            diststyle_auto = re.search(r'DISTSTYLE\s+AUTO', create_table_stmt, flags=re.IGNORECASE)
            # 检查是否在列定义或表属性中指定了 DISTKEY
            distkey_specified = re.search(r'\bDISTKEY\b', create_table_stmt, flags=re.IGNORECASE)

            if diststyle_auto and distkey_specified:
                # 如果同时指定了 DISTSTYLE AUTO 和 DISTKEY，将 DISTSTYLE AUTO 更改为 DISTSTYLE KEY
                create_table_stmt = re.sub(
                    r'DISTSTYLE\s+AUTO',
                    'DISTSTYLE KEY',
                    create_table_stmt,
                    flags=re.IGNORECASE
                )
            elif diststyle_auto and not distkey_specified:
                # 如果指定了 DISTSTYLE AUTO，但未指定 DISTKEY，不做任何修改
                pass
            elif not diststyle_auto and distkey_specified:
                # 如果未指定 DISTSTYLE，但指定了 DISTKEY，不做任何修改
                pass
            else:
                # 如果未指定 DISTSTYLE，且未指定 DISTKEY，可以根据需要设置默认的 DISTSTYLE
                pass

            # 移除多余的分号和空白
            create_table_stmt = create_table_stmt.strip()
            # 检查语句结尾是否已有分号
            if not create_table_stmt.endswith(';'):
                create_table_stmt += ';'

            ddl_statements.append(create_table_stmt + "\n")
            table_count += 1

            # 每处理50个表，保存到一个文件
            if table_count % 50 == 0:
                filename = f"dev_yujiahao_create_table_test{'' if file_count == 0 else file_count}.sql"
                with open(filename, 'w', encoding='utf-8') as f:
                    f.write('\n'.join(ddl_statements))
                start_table = file_count * 50 + 1
                end_table = start_table + 49
                print(f"已生成 {filename}，包含第 {start_table} 到 {end_table} 个表的 CREATE TABLE 语句")
                # 清空列表，为下一个文件做准备
                ddl_statements = []
                file_count += 1

        except Exception as e:
            print(f"处理表 {table_name} 时出错: {e}")
            continue  # 跳过当前表，继续处理下一个表

    # 处理少于50个的剩余表
    if ddl_statements:
        filename = f"dev_yujiahao_create_table_test{'' if file_count == 0 else file_count}.sql"
        with open(filename, 'w', encoding='utf-8') as f:
            f.write('\n'.join(ddl_statements))
        start_table = file_count * 50 + 1
        end_table = start_table + len(ddl_statements) - 1
        print(f"已生成 {filename}，包含第 {start_table} 到 {end_table} 个表的 CREATE TABLE 语句")

    # 关闭游标和连接
    cur.close()
    conn.close()


if __name__ == '__main__':
    main()
